package cn.miao.middle.cloud.core.async;

import java.util.Map;

import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

/**
 * 传递RequestAttributes and MDC and SecurityContext
 * 
 * @author kongyufei
 *
 */
public class ContextCopyingDecorator implements TaskDecorator {
	@Override
	public Runnable decorate(Runnable runnable) {
		try {
			RequestAttributes context = RequestContextHolder.currentRequestAttributes();
			Map<String, String> previous = MDC.getCopyOfContextMap();
			SecurityContext securityContext = SecurityContextHolder.getContext();
			return () -> {
				try {
					RequestContextHolder.setRequestAttributes(context);
					MDC.setContextMap(previous);
					SecurityContextHolder.setContext(securityContext);
					runnable.run();
				} finally {
					RequestContextHolder.resetRequestAttributes();
					MDC.clear();
					SecurityContextHolder.clearContext();
				}
			};
		} catch (IllegalStateException e) {
			return runnable;
		}
	}
}
